package com.liuqi.openai.model;

import com.liuqi.openai.core.client.OpenAiRemoteException;
import com.liuqi.openai.core.embedding.EmbeddingRequest;
import com.liuqi.openai.core.embedding.EmbeddingResponse;
import com.liuqi.openai.data.embedding.Embedding;
import com.liuqi.openai.model.common.AbstractOpenAiModel;
import com.liuqi.openai.model.embedding.EmbeddingModel;
import com.liuqi.openai.model.output.Response;
import java.util.List;

import static java.util.stream.Collectors.toList;


/**
 * @author liuqi
 * @date 2025/7/20
 **/
public class OpenAiEmbeddingModel extends AbstractOpenAiModel implements EmbeddingModel {

    private final Integer dimensions;

    OpenAiEmbeddingModel(OpenAiEmbeddingModelBuilder builder) {
        super(builder);
        this.dimensions = builder.dimensions;
    }

    public static OpenAiEmbeddingModelBuilder builder() {
        return new OpenAiEmbeddingModelBuilder();
    }

    @Override
    public Response<List<Embedding>> embed(List<String> texts) throws OpenAiModelException {
        if (texts == null || texts.isEmpty()) {
            throw new OpenAiModelException("texts cannot be null or empty");
        }

        try {
            // 构建 Embedding 入参
            EmbeddingRequest request = EmbeddingRequest.builder()
                    .model(modelName)
                    .input(texts)
                    .dimensions(dimensions)
                    .build();

            // 发起 embedding 请求
            EmbeddingResponse response = openAiClient.embedding(request).execute();

            // 构建响应内容
            List<Embedding> embeddings = response.getData().stream()
                    .map(embedding -> Embedding.from(embedding.getEmbedding()))
                    .collect(toList());

            // 构建 Response
            return Response.from(
                    embeddings,
                    tokenUsageFrom(response.getUsage())
            );
        } catch (OpenAiRemoteException e) {
            // 远程服务异常
            throw new OpenAiModelException("OpenAiEmbeddingModel Error.", e);
        }
    }

    public static class OpenAiEmbeddingModelBuilder
            extends AbstractOpenAiModelBuilder<OpenAiEmbeddingModel, OpenAiEmbeddingModelBuilder> {
        Integer dimensions;

        public OpenAiEmbeddingModelBuilder dimensions(Integer dimensions) {
            this.dimensions = dimensions;
            return this;
        }

        @Override
        public OpenAiEmbeddingModel build() {
            return new OpenAiEmbeddingModel(this);
        }
    }

}
